package org.shuzhou.d_stateful;

import org.shuzhou.model.PersonalInfo;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.functions.RichReduceFunction;
import org.apache.flink.api.common.serialization.SimpleStringSchema;
import org.apache.flink.api.common.state.MapState;
import org.apache.flink.api.common.state.MapStateDescriptor;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.connector.kafka.source.KafkaSource;
import org.apache.flink.connector.kafka.source.enumerator.initializer.OffsetsInitializer;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.api.common.typeinfo.Types;

/**
 * 数据集：data/people.txt
 * 计算每个民族的人数，使用MapState存储中间结果
 */
public class PersonCountStateful {
    public static void main(String[] args) throws Exception {

        // set up the streaming execution environment
        final StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        /**
         * 获取数据，数据源为Kafka，避免从端口获取中文数据乱码
         * 数据格式
         * 王矗馨,女,汉族,411424195001010028,19500101,150977754190,kk4jlw850x@sina.com,甘肃省临夏回族自治州和政县CNR路884号
         */
        KafkaSource<String> source = KafkaSource.<String>builder()
                .setBootstrapServers("192.168.56.151:9092,192.168.56.152:9092,192.168.56.153:9092")
                .setTopics("persion")
                .setGroupId("root")
                .setStartingOffsets(OffsetsInitializer.earliest())
                .setValueOnlyDeserializer(new SimpleStringSchema())
                .build();

        DataStream<String> people = env.fromSource(source, WatermarkStrategy.noWatermarks(), "Kafka Source");

        /**
         * 获取数据，数据源为kafka
         * 数据格式
         * 王矗馨,女,汉族,411424195001010028,19500101,150977754190,kk4jlw850x@sina.com,甘肃省临夏回族自治州和政县CNR路884号
         * 计算后返回Tuple2（民族，PersonalInfo对象）
         */
        SingleOutputStreamOperator<Tuple2<String, PersonalInfo>> ds = people.map(new MapFunction<String, Tuple2<String, PersonalInfo>>() {
            @Override
            public Tuple2<String, PersonalInfo> map(String value) throws Exception {
                // System.out.println(value.toString());
                String encode = new String(value.getBytes("UTF-8"), "UTF-8");
                String[] fields = encode.split(",");
                PersonalInfo info = new PersonalInfo(fields[0], fields[1], fields[2], fields[3], fields[4], fields[5],
                        fields[6], fields[7]);
                return Tuple2.of(fields[2],info);
            }
        });

        /**
         * 使用keyBy对民族进行分组
         * reduce聚合案例：对每个民族的人数与年龄进行累加
         * 使用有状态算子 keyed state
         */
        ds.map(new MapFunction<Tuple2<String, PersonalInfo>, Tuple2<String, Integer>>() {
            @Override
            public Tuple2<String, Integer> map(Tuple2<String, PersonalInfo> value) throws Exception {
                return Tuple2.of(value.f0,1);
            }
        })
            .keyBy(value -> value.f0)
            .reduce(new RichReduceFunction<Tuple2<String, Integer>>() {
            // sum用于记录状态，f0=民族，f1=人口总数
            private transient MapState<String, Integer> sum;

            @Override
            public Tuple2<String, Integer> reduce(Tuple2<String, Integer> value1, Tuple2<String, Integer> value2) throws Exception {
                
                // 获取sum状态
                Integer currentSum = sum.get(value1.f0);

                if (currentSum == null) {
                    currentSum = 1;
                }

                if (!value2.equals(null)) {
                    // 人数+1
                    currentSum += 1;
                }

                // 将民族，人数写回map
                sum.put(value1.f0, currentSum);

                // 将计算结果返回 (民族，人数）
                return Tuple2.of(value1.f0, currentSum);
            }

            // 源码提示中推荐的写法
            @Override
            public void open(Configuration config) throws Exception{
                // 构造时不初始化默认值
                MapStateDescriptor<String, Integer> descriptor = new MapStateDescriptor("sum", Types.STRING, Types.INT);
                sum = getRuntimeContext().getMapState(descriptor);
                // 在下面的代码中，判断如果state没有初始化则赋为默认值
                if (sum == null) {
                    sum.put("null", 0);
                }
            }

        }).print();

        env.execute("Flink Operator Example");
    }
}
